package com.passion.filter;

import org.owasp.validator.html.Policy;
import org.owasp.validator.html.PolicyException;
import org.owasp.validator.html.ScanException;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * 类描述：Xss过滤器
 *
 * @author ZhuYin
 * @since 2024年12月26日
 */
@Component
@Order(1)
public class XssFilter extends OncePerRequestFilter {

    private static Policy policy;

    public XssFilter() throws IOException, PolicyException {
        // 加载 AntiSamy 策略文件
        XssFilter.policy = Policy.getInstance(new ClassPathResource("antisamy-anythinggoes.xml").getInputStream());
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain) throws ServletException, IOException {
        // 检查是否是静态资源或不需要过滤的路径
        if (isExcludedPath(request)) {
            filterChain.doFilter(request, response);
            return;
        }

        // // 检查是否是JSON请求
        // String contentType = request.getContentType();
        // if (contentType != null && contentType.contains("application/json")) {
        //     // 对于JSON请求，直接传递原始请求
        //     filterChain.doFilter(request, response);
        //     return;
        // }

        // 创建带有XSS清理功能的包装请求对象
        XssHttpServletRequestWrapper xssRequest = new XssHttpServletRequestWrapper(request);
        filterChain.doFilter(xssRequest, response);
    }

    private boolean isExcludedPath(HttpServletRequest request) {
        String uri = request.getRequestURI();
        // 排除静态资源、登录页面等不需要过滤的路径
        return uri.contains("/css/") || uri.contains("/js/") || uri.contains("/images/") || uri.contains("/webjars/");
    }

    private static class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

        private final Map<String, String[]> cleanedParameters;

        public XssHttpServletRequestWrapper(HttpServletRequest request) {
            super(request);
            cleanedParameters = cleanParameters(request.getParameterMap());
        }

        @Override
        public String[] getParameterValues(String parameter) {
            String[] values = cleanedParameters.get(parameter);
            if (values == null) {
                return super.getParameterValues(parameter);
            }
            return values;
        }

        @Override
        public String getParameter(String parameter) {
            String[] values = getParameterValues(parameter);
            if (values != null && values.length > 0) {
                return values[0];
            }
            return null;
        }

        @Override
        public Map<String, String[]> getParameterMap() {
            return cleanedParameters;
        }

        private Map<String, String[]> cleanParameters(Map<String, String[]> parameters) {
            Map<String, String[]> cleanedParams = new HashMap<>();
            for (Map.Entry<String, String[]> entry : parameters.entrySet()) {
                String[] cleanedValues = new String[entry.getValue().length];
                for (int i = 0; i < entry.getValue().length; i++) {
                    try {
                        cleanedValues[i] = sanitize(entry.getValue()[i]);
                    } catch (ScanException | PolicyException e) {
                        // 如果清理失败，可以记录日志或设置为空字符串
                        cleanedValues[i] = "";
                    }
                }
                cleanedParams.put(entry.getKey(), cleanedValues);
            }
            return cleanedParams;
        }

        private String sanitize(String value) throws ScanException, PolicyException {
            if (value == null || value.isEmpty()) {
                return value;
            }
            // 使用 AntiSamy 进行清理
            return new org.owasp.validator.html.AntiSamy().scan(value, policy).getCleanHTML();
        }
    }
}